import sys
import cv2
import numpy as np
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout,
                             QLabel, QPushButton, QFileDialog, QSlider, QSpinBox, QComboBox,
                             QTextEdit, QTabWidget, QGroupBox, QSplitter, QMessageBox)
from PyQt5.QtGui import QImage, QPixmap, QPainter, QPen, QColor
from PyQt5.QtCore import Qt, QThread, pyqtSignal, QPoint
import os
from image_processing import ImageProcessor
from computer_vision import FaceRecognizer
from video_processing import VideoProcessor


class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("图像处理与计算机视觉系统")
        self.setGeometry(100, 100, 1200, 800)

        # 初始化处理器
        self.image_processor = ImageProcessor()
        self.face_recognizer = FaceRecognizer()
        self.video_processor = VideoProcessor()

        # 存储当前图像和处理历史
        self.current_image = None
        self.original_image = None
        self.processing_history = []

        # 创建UI
        self.init_ui()

        # 连接信号和槽
        self.connect_signals()

    def init_ui(self):
        # 创建主布局
        main_widget = QWidget()
        main_layout = QHBoxLayout(main_widget)
        self.setCentralWidget(main_widget)

        # 左侧控制面板
        control_panel = QVBoxLayout()

        # 文件操作组
        file_group = QGroupBox("文件操作")
        file_layout = QVBoxLayout()

        self.open_image_btn = QPushButton("打开图像")
        self.save_image_btn = QPushButton("保存图像")
        self.open_video_btn = QPushButton("打开视频")
        self.start_camera_btn = QPushButton("启动摄像头")

        file_layout.addWidget(self.open_image_btn)
        file_layout.addWidget(self.save_image_btn)
        file_layout.addWidget(self.open_video_btn)
        file_layout.addWidget(self.start_camera_btn)

        file_group.setLayout(file_layout)
        control_panel.addWidget(file_group)

        # 图像处理组
        process_group = QGroupBox("图像处理")
        process_layout = QVBoxLayout()

        self.process_type = QComboBox()
        self.process_type.addItems([
            "无操作", "旋转", "缩放", "裁剪",
            "亮度/对比度", "饱和度", "平滑处理",
            "形状识别", "图像拼接", "人脸识别"
        ])

        self.process_btn = QPushButton("执行处理")

        process_layout.addWidget(self.process_type)
        process_layout.addWidget(self.process_btn)

        # 参数区域
        self.param_area = QWidget()
        self.param_layout = QVBoxLayout(self.param_area)
        process_layout.addWidget(self.param_area)

        process_group.setLayout(process_layout)
        control_panel.addWidget(process_group)

        # 视频控制组
        video_group = QGroupBox("视频控制")
        video_layout = QVBoxLayout()

        self.play_btn = QPushButton("播放")
        self.pause_btn = QPushButton("暂停")
        self.stop_btn = QPushButton("停止")
        self.forward_btn = QPushButton("快进")
        self.backward_btn = QPushButton("快退")

        video_layout.addWidget(self.play_btn)
        video_layout.addWidget(self.pause_btn)
        video_layout.addWidget(self.stop_btn)
        video_layout.addWidget(self.forward_btn)
        video_layout.addWidget(self.backward_btn)

        video_group.setLayout(video_layout)
        control_panel.addWidget(video_group)

        control_panel.addStretch(1)
        main_layout.addLayout(control_panel, 1)

        # 右侧显示区域
        display_splitter = QSplitter(Qt.Vertical)

        # 图像/视频显示区域
        self.display_area = QLabel("请打开图像或视频")
        self.display_area.setAlignment(Qt.AlignCenter)
        self.display_area.setMinimumSize(400, 300)
        self.display_area.setStyleSheet("border: 1px solid #ccc;")
        self.display_area.setScaledContents(True)
        display_splitter.addWidget(self.display_area)

        # 结果显示区域
        self.result_area = QTextEdit()
        self.result_area.setReadOnly(True)
        self.result_area.setMinimumHeight(200)
        display_splitter.addWidget(self.result_area)

        main_layout.addWidget(display_splitter, 3)

        # 初始化参数区域
        self.init_param_area()

    def init_param_area(self):
        # 清空参数区域
        while self.param_layout.count():
            item = self.param_layout.takeAt(0)
            widget = item.widget()
            if widget:
                widget.deleteLater()

        # 根据当前处理类型显示不同参数
        process_type = self.process_type.currentText()

        if process_type == "旋转":
            angle_layout = QHBoxLayout()
            angle_label = QLabel("旋转角度:")
            self.angle_spinbox = QSpinBox()
            self.angle_spinbox.setRange(-360, 360)
            self.angle_spinbox.setValue(0)
            angle_layout.addWidget(angle_label)
            angle_layout.addWidget(self.angle_spinbox)
            self.param_layout.addLayout(angle_layout)

        elif process_type == "缩放":
            scale_layout = QHBoxLayout()
            scale_label = QLabel("缩放比例(%):")
            self.scale_spinbox = QSpinBox()
            self.scale_spinbox.setRange(10, 500)
            self.scale_spinbox.setValue(100)
            scale_layout.addWidget(scale_label)
            scale_layout.addWidget(self.scale_spinbox)
            self.param_layout.addLayout(scale_layout)

        elif process_type == "亮度/对比度":
            brightness_layout = QHBoxLayout()
            brightness_label = QLabel("亮度:")
            self.brightness_slider = QSlider(Qt.Horizontal)
            self.brightness_slider.setRange(-100, 100)
            self.brightness_slider.setValue(0)
            brightness_layout.addWidget(brightness_label)
            brightness_layout.addWidget(self.brightness_slider)
            self.param_layout.addLayout(brightness_layout)

            contrast_layout = QHBoxLayout()
            contrast_label = QLabel("对比度:")
            self.contrast_slider = QSlider(Qt.Horizontal)
            self.contrast_slider.setRange(-100, 100)
            self.contrast_slider.setValue(0)
            contrast_layout.addWidget(contrast_label)
            contrast_layout.addWidget(self.contrast_slider)
            self.param_layout.addLayout(contrast_layout)

        elif process_type == "饱和度":
            saturation_layout = QHBoxLayout()
            saturation_label = QLabel("饱和度:")
            self.saturation_slider = QSlider(Qt.Horizontal)
            self.saturation_slider.setRange(-100, 100)
            self.saturation_slider.setValue(0)
            saturation_layout.addWidget(saturation_label)
            saturation_layout.addWidget(self.saturation_slider)
            self.param_layout.addLayout(saturation_layout)

        elif process_type == "平滑处理":
            smooth_layout = QHBoxLayout()
            smooth_label = QLabel("平滑类型:")
            self.smooth_combo = QComboBox()
            self.smooth_combo.addItems(["均值滤波", "高斯滤波"])
            smooth_layout.addWidget(smooth_label)
            smooth_layout.addWidget(self.smooth_combo)
            self.param_layout.addLayout(smooth_layout)

            kernel_layout = QHBoxLayout()
            kernel_label = QLabel("核大小:")
            self.kernel_spinbox = QSpinBox()
            self.kernel_spinbox.setRange(1, 21)
            self.kernel_spinbox.setValue(3)
            self.kernel_spinbox.setSingleStep(2)
            kernel_layout.addWidget(kernel_label)
            kernel_layout.addWidget(self.kernel_spinbox)
            self.param_layout.addLayout(kernel_layout)

        elif process_type == "图像拼接":
            self.add_image_btn = QPushButton("添加待拼接图像")
            self.param_layout.addWidget(self.add_image_btn)

            self.images_to_stitch = []

        elif process_type == "人脸识别":
            self.register_face_btn = QPushButton("注册人脸")
            self.param_layout.addWidget(self.register_face_btn)

    def connect_signals(self):
        # 文件操作
        self.open_image_btn.clicked.connect(self.open_image)
        self.save_image_btn.clicked.connect(self.save_image)
        self.open_video_btn.clicked.connect(self.open_video)
        self.start_camera_btn.clicked.connect(self.start_camera)

        # 处理类型变化
        self.process_type.currentIndexChanged.connect(self.init_param_area)

        # 处理按钮
        self.process_btn.clicked.connect(self.process_image)

        # 视频控制
        self.play_btn.clicked.connect(self.play_video)
        self.pause_btn.clicked.connect(self.pause_video)
        self.stop_btn.clicked.connect(self.stop_video)
        self.forward_btn.clicked.connect(self.forward_video)
        self.backward_btn.clicked.connect(self.backward_video)

        # 图像拼接
        if hasattr(self, 'add_image_btn'):
            self.add_image_btn.clicked.connect(self.add_stitch_image)

        # 人脸识别
        if hasattr(self, 'register_face_btn'):
            self.register_face_btn.clicked.connect(self.register_face)

        # 视频处理器信号
        self.video_processor.frame_ready.connect(self.update_video_frame)
        self.video_processor.processing_result.connect(self.update_processing_result)

    def open_image(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "打开图像", "", "图像文件 (*.png *.jpg *.jpeg *.bmp)"
        )
        if file_path:
            self.current_image = cv2.imread(file_path)
            self.original_image = self.current_image.copy()
            self.display_image(self.current_image)
            self.result_area.append(f"已加载图像: {os.path.basename(file_path)}")

    def save_image(self):
        if self.current_image is not None:
            file_path, _ = QFileDialog.getSaveFileName(
                self, "保存图像", "", "图像文件 (*.png *.jpg *.jpeg *.bmp)"
            )
            if file_path:
                cv2.imwrite(file_path, self.current_image)
                self.result_area.append(f"已保存图像: {os.path.basename(file_path)}")
        else:
            QMessageBox.warning(self, "警告", "没有可保存的图像")

    def open_video(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "打开视频", "", "视频文件 (*.mp4 *.avi *.mov *.mkv)"
        )
        if file_path:
            self.video_processor.open_video(file_path)
            self.result_area.append(f"已加载视频: {os.path.basename(file_path)}")

    def start_camera(self):
        self.video_processor.open_camera()
        self.result_area.append("已启动摄像头")

    def process_image(self):
        if self.current_image is None:
            QMessageBox.warning(self, "警告", "请先加载图像")
            return

        process_type = self.process_type.currentText()

        try:
            if process_type == "旋转":
                angle = self.angle_spinbox.value()
                self.current_image = self.image_processor.rotate(self.current_image, angle)
                self.result_area.append(f"已旋转图像: {angle}度")

            elif process_type == "缩放":
                scale = self.scale_spinbox.value() / 100.0
                self.current_image = self.image_processor.scale(self.current_image, scale)
                self.result_area.append(f"已缩放图像: {scale * 100}%")

            elif process_type == "裁剪":
                # 这里需要实现裁剪区域选择
                self.result_area.append("裁剪功能需要交互选择区域，将在后续版本实现")

            elif process_type == "亮度/对比度":
                brightness = self.brightness_slider.value()
                contrast = self.contrast_slider.value()
                self.current_image = self.image_processor.adjust_brightness_contrast(
                    self.current_image, brightness, contrast
                )
                self.result_area.append(f"已调整亮度: {brightness}, 对比度: {contrast}")

            elif process_type == "饱和度":
                saturation = self.saturation_slider.value()
                self.current_image = self.image_processor.adjust_saturation(
                    self.current_image, saturation
                )
                self.result_area.append(f"已调整饱和度: {saturation}")

            elif process_type == "平滑处理":
                smooth_type = self.smooth_combo.currentText()
                kernel_size = self.kernel_spinbox.value()
                if smooth_type == "均值滤波":
                    self.current_image = self.image_processor.mean_filter(
                        self.current_image, kernel_size
                    )
                elif smooth_type == "高斯滤波":
                    self.current_image = self.image_processor.gaussian_filter(
                        self.current_image, kernel_size
                    )
                self.result_area.append(f"已进行{smooth_type}: 核大小{kernel_size}")

            elif process_type == "形状识别":
                shapes, result_image = self.image_processor.detect_shapes(self.current_image)
                self.current_image = result_image
                self.result_area.append(f"已识别形状: {len(shapes)}个")
                for i, shape in enumerate(shapes):
                    self.result_area.append(f"形状 {i + 1}: {shape['type']}, 位置: {shape['center']}")

            elif process_type == "图像拼接":
                if hasattr(self, 'images_to_stitch') and len(self.images_to_stitch) > 0:
                    images = [self.current_image] + self.images_to_stitch
                    self.current_image = self.image_processor.stitch_images(images)
                    self.result_area.append(f"已拼接{len(images)}张图像")
                    # 清空待拼接图像列表
                    self.images_to_stitch = []
                    # 更新参数区域
                    self.init_param_area()
                else:
                    self.result_area.append("请先添加待拼接图像")

            elif process_type == "人脸识别":
                faces, result_image = self.face_recognizer.detect_faces(self.current_image)
                self.current_image = result_image
                self.result_area.append(f"已检测到{len(faces)}个人脸")
                for i, face in enumerate(faces):
                    if 'name' in face:
                        self.result_area.append(f"人脸 {i + 1}: {face['name']}, 相似度: {face['confidence']:.2f}")
                    else:
                        self.result_area.append(f"人脸 {i + 1}: 未识别")

            # 显示处理后的图像
            self.display_image(self.current_image)

        except Exception as e:
            QMessageBox.critical(self, "错误", f"处理图像时出错: {str(e)}")
            self.result_area.append(f"错误: {str(e)}")

    def add_stitch_image(self):
        file_path, _ = QFileDialog.getOpenFileName(
            self, "添加待拼接图像", "", "图像文件 (*.png *.jpg *.jpeg *.bmp)"
        )
        if file_path:
            image = cv2.imread(file_path)
            self.images_to_stitch.append(image)
            self.result_area.append(f"已添加待拼接图像: {os.path.basename(file_path)}")

    def register_face(self):
        if self.current_image is None:
            QMessageBox.warning(self, "警告", "请先加载包含人脸的图像")
            return

        name, ok = QMessageBox.getText(self, "注册人脸", "请输入此人姓名:")
        if ok and name:
            faces, _ = self.face_recognizer.detect_faces(self.current_image)
            if len(faces) > 0:
                # 假设第一个检测到的人脸是要注册的
                face = faces[0]
                success = self.face_recognizer.register_face(
                    self.current_image, face['box'], name
                )
                if success:
                    self.result_area.append(f"已注册人脸: {name}")
                else:
                    self.result_area.append("注册人脸失败")
            else:
                self.result_area.append("未检测到人脸")

    def play_video(self):
        self.video_processor.play()
        process_type = self.process_type.currentText()
        if process_type == "人脸识别":
            self.video_processor.set_processing_function(self.face_recognizer.detect_faces)
        else:
            self.video_processor.set_processing_function(None)

    def pause_video(self):
        self.video_processor.pause()

    def stop_video(self):
        self.video_processor.stop()

    def forward_video(self):
        self.video_processor.forward()

    def backward_video(self):
        self.video_processor.backward()

    def display_image(self, image):
        if image is not None:
            # 转换为RGB格式
            rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            h, w, ch = rgb_image.shape
            bytes_per_line = ch * w
            q_image = QImage(rgb_image.data, w, h, bytes_per_line, QImage.Format_RGB888)
            self.display_area.setPixmap(QPixmap.fromImage(q_image))

    def update_video_frame(self, frame):
        self.display_image(frame)

    def update_processing_result(self, result):
        if result:
            self.result_area.append(result)